#!/usr/bin/env python

# --------
# Imports:
# --------
# Project configuration 
from config import *

# Ros imports:
import roslib; roslib.load_manifest(NODE_NAME)
import rospy
from geometry_msgs.msg import Pose
from nav_msgs.msg import OccupancyGrid
from sensor_msgs.msg import CompressedImage

# Python imports
from copy import deepcopy
from math import sin, cos, sqrt, exp, pi, log, pow
from random import gauss, sample
from StringIO import StringIO
from time import time
import Image

# Other imports
from misc import *
from bresenham_line import *

# ----------------
# Global variables 
# ----------------
particleId = -1 

# ---------------------
def generateUniqueID():
# ---------------------
	global particleId
	particleId += 1
	return particleId

# -------------
class Particle:
# -------------
	# ---------------------------------
	def __init__(self, parent, slamMap):
	# ---------------------------------
		# Members:
		# -------
		self.position = None						# The particle's position (world coordinates)
		self.heading = None							# The particle's heading (world coordinates)
		self.particleId = generateUniqueID() 		# The particle's unique ID
		self.parent = parent						# The particle's parent
		self.numChildren = 0						# We haven't been sampled yet, so no children
		self.probability = 1						# Probability defaults to one
		self.slamMap = slamMap						# A reference to the map data structure
		self.updated = set()						# The set of this particle's updated map cells
	
	# --------------
	def sample(self):
	# --------------
		""" Copy the particle for sampling  """
		
		# Create a new particle with this particle as its parent
		sampled = Particle(self, self.slamMap)

		# Copy the current heading
		sampled.heading = deepcopy(self.heading)

		# Copy the current position
		sampled.position = deepcopy(self.position)

		# The particle has now been sampled!
		self.numChildren += 1

		# Return the sampled particle
		return sampled

	# ------------------------------
	def localize(self, laserScanMsg):
	# ------------------------------
		""" Run localization at this particle's position in this particle's map """
		
		# Get the most recent position/heading
		if self.position is None: return 

		position = self.position
		heading = self.heading

		# Initializations
		mapRanges = []
		reconstructedCells = {}
		ancestryList = self.getAncestryList()

		# Cache some reused elements as local variables
		angleMin = laserScanMsg.angle_min
		angleIncrement = laserScanMsg.angle_increment
		mapInBounds = self.slamMap.mapInBounds
		slamMap = self.slamMap
		mapRangesAppend = mapRanges.append
		getMapCell = self.getMapCell

		# Generate LASER_LOCALIZE_NUM_SAMPLE random indices (for randomly sampling out lasers), filter out any garbage, and calculate bearings
		randomLaserIndices = sample(range(len(laserScanMsg.ranges)), LASER_LOCALIZE_NUM_SAMPLE)
		distInRange = lambda x: x[0] >= laserScanMsg.range_min and x[0] <= laserScanMsg.range_max 
		calculateBearing = lambda i: clipAngle((i * angleIncrement) + angleMin + heading)
		calcRayEnd = lambda bearing: self.slamMap.worldToMap(((position[0] + cos(bearing) * MAP_MAX_LOCALIZE_RANGE_M), position[1] + (sin(bearing) * MAP_MAX_LOCALIZE_RANGE_M)))
		laserData = filter(distInRange, [(laserScanMsg.ranges[i], calcRayEnd(calculateBearing(i))) for i in randomLaserIndices])

		rayStart = self.slamMap.worldToMap(position)
		for distance, rayEnd in laserData:

			# Iterate through rayPts and reconstruct map value for each point, stopping and storing the distance when reach an obstacle
			obstacle = None
			for pt in bresenham_line_gen(rayStart, rayEnd):

				# Make sure the point is in bounds
				if mapInBounds(pt):
					# Get the cell value
					if not pt in reconstructedCells: 
						reconstructedCells[pt] = getMapCell(pt, ancestryList)
					cellValue = reconstructedCells[pt]

					# Break if cellValue is MAP_CELL_ON
					if cellValue == MAP_CELL_ON:
						obstacle = pt
						break
					elif cellValue == MAP_CELL_UNKNOWN:
						break
				
			# If didn't encounter obstacle, store the maximum range in mapRanges
			if obstacle == None:
				mapRangesAppend(MAP_MAX_LOCALIZE_RANGE_M)

			# Otherwise, store the distance to obstacle
			else:
				dx = (obstacle[0]-rayStart[0])
				dy = (obstacle[1]-rayStart[1])
				distance = sqrt(dx*dx + dy*dy) * MAP_METERS_PER_PIXEL
				mapRangesAppend(distance)
		
		# Determine probability of seeing the the map ranges given the actual laser ranges we saw
		self.poseProbability(mapRanges, [ld[0] for ld in laserData])
		
	# -----------------------------------------------
	def poseProbability(self, mapRanges, laserRanges):
	# -----------------------------------------------
		""" Given a list of map ranges and laser ranges, determine the probability of being in the particle's pose """
		probability = 0

		# For each pair of laser and map ranges, find the probability of getting the sensor reading
		for i in range(0, len(mapRanges)):
			slamMap = mapRanges[i]

			# Ignore places where we weren't able to raytrace out to an obstacle in our map
			if slamMap != MAP_MAX_LOCALIZE_RANGE_M:
				laser = laserRanges[i]
				#prob = (gaussian(slamMap, SENSOR_SIGMA, laser) / gaussian(slamMap, SENSOR_SIGMA, slamMap))
				prob = gaussian(slamMap, SENSOR_SIGMA, laser) 

				# Taking the log of 0 is undefined , so don't do it
				probability += log(max(prob,0.0000000000000000001))

		# Multiplying such tiny floating point numbers isn't reliable, so we're adding logs instead
		self.probability = exp(probability)

	# ------------------------------
	def updateMap(self, laserScanMsg):
	# ------------------------------
	# Given a laser scan message (in the robot's reference frame), update the particle's map
	
		# Cache commonly used functions as local variables
		worldToMap = self.slamMap.worldToMap
		setInMap = self.setInMap
	
		# Get the most recent pose (last in the list)
		if self.position is None: return
		position = self.position
		heading  = self.heading

		# Store all the points swept out by the laser rangefinder and all the contact pts
		sweptPts = set()
		contactPts = set()

		# Loop through the laser scan
		for i, r in enumerate(laserScanMsg.ranges):

			# Our laser reads max range as 0...
			if r == 0.0: r = MAP_MAX_UPDATE_RANGE_M

			# Ignore garbage readings
			if (r >= laserScanMsg.range_min and r <= laserScanMsg.range_max):
				
				# Calculate the bearing of this range
				bearing = clipAngle((i * laserScanMsg.angle_increment) + laserScanMsg.angle_min + heading)

				# Get the ray's start point and end point (and convert them to map coordinates)
				rayStart = (position[0], position[1])
				rayEnd = (rayStart[0] + (cos(bearing) * r), rayStart[1] + (sin(bearing) * r))
				rayStart = worldToMap(rayStart)
				rayEnd = worldToMap(rayEnd)

				# Run Bresenham's algorithm to get a list of all the points between rayStart and rayEnd
				rayPts = bresenham_line(rayStart, rayEnd)
			
				# Add the last ray point to the set of contact points	
				if r != MAP_MAX_UPDATE_RANGE_M and self.slamMap.mapInBounds(rayPts[-1]):
					contactPts.add(rayPts[-1])
			
				# Filter the swept points to make sure they're all in bounds in our map data structure
				rayPts = filter(self.slamMap.mapInBounds, rayPts)

				# Update the set of swept out points
				sweptPts |= set(rayPts)

		# Remove any contact points from the swept point list
		sweptPts -= (sweptPts & contactPts)
		
		# Get our ancestry chain
		ancestryList = self.getAncestryList()

		# Now go through all of our swept points and set their value to 0 in our map
		for pt in sweptPts:
			setInMap(pt, MAP_CELL_OFF, ancestryList)

		# Finally, go through all of our contact points and set their values to 1 in our map
		for pt in contactPts:
			setInMap(pt, MAP_CELL_ON, ancestryList)

	# -----------------------
	def getAncestryList(self):
	# -----------------------
		# Initialize the container we're going to use
		ancestry = []
	
		# Iteratively walk up the ancestry chain until you walk past the root node
		pointer = self
		while pointer is not None:
			ancestry.append(pointer.particleId)
			pointer = pointer.parent

		# Return the ancestry chain
		return ancestry
	
	# ---------------------------------------
	def setInMap(self, pt, val, ancestryList):
	# ---------------------------------------
		# Alias the map cell at pt
		cell = self.slamMap[pt]

		# Lock the map cell we're examining
		with self.slamMap.getLock(pt):

			# Walk along the ancestry chain
			for ancestor in ancestryList:

				# See if this ancestor touched the value of the cell at pt
				if ancestor in cell:
					
					# See if the ancestor *changed* the value of the cell at pt
					if cell[ancestor] != val:
					
						# Add this point as a point we've changed from our ancestor
						self.updated.add(pt)

						# Indicate in the grid cell that this particle is changing the value
						cell[self.particleId] = val 
					
					# In either case, we're done with this function now
					return None
		
			# If we get here then none of our ancestors have seen the particle
			self.updated.add(pt)
			cell[self.particleId] = val

	# ------------------------------
	def generateOccupancyGrid(self):
	# ------------------------------
		# Create the occupancy grid
		grid = OccupancyGrid()

		# Set the map meta-data
		grid.info.resolution = self.slamMap.metersPerPixel
		grid.info.width = self.slamMap.widthPx
		grid.info.height = self.slamMap.heightPx
		grid.info.origin = Pose() # TO DO: Fix this!

		# Get the ancestry chain for this particle
		ancestryList = self.getAncestryList()

		# Loop through all of the map cells: (in row-major order)       
		for y in range(self.slamMap.heightPx):
			for x in range(self.slamMap.widthPx):

				# Reconstruct the map cell at (x,y)
				cell = self.slamMap[(x,y)]

				# The cell value defaults to unknown
				cellValue = MAP_CELL_UNKNOWN

				# Look for the first ancestor to have changed the value
				for ancestor in ancestryList:
					# If the ancestor changed the value of the cell, we're done
					if ancestor in cell:
						cellValue = cell[ancestor]
						break

				# Set the cell value in the data grid
				grid.data.append(cellValue)

		return grid

	# -----------------------------
	def generateVisualization(self):
	# -----------------------------
		# Get the occupancy grid and alias the width and height
		occupancyGrid = self.generateOccupancyGrid()
		width = occupancyGrid.info.width
		height = occupancyGrid.info.height

		# Use the python image library to create a new image
		img = Image.new('L', (width, height))

		# Loop through the occupancy grid data (which is stored in row-major order)
		for i, val in enumerate(occupancyGrid.data):
			# Calculate the point
			pt = (i % width, i / width)

			if val == MAP_CELL_UNKNOWN:
				img.putpixel(pt, 127)
			else:	
				# Put the pixel	into the image
				img.putpixel(pt, int((1-(val/100.0)) * 255))

		# Return the image!		
		return img

	# -------------
	def prune(self):
	# -------------
	# Recursively prune the branch of the tree with this particle as a leaf and remove
	# all corresponding entries into the map grid structure
		if self.numChildren == 0:
			
			# Remove the particle from each of the map grid cells that it updated
			for pt in self.updated:
				cell = self.slamMap[pt]
				del cell[self.particleId]
			
			# Recurse!
			if self.parent is not None:	
				# The parent now has one less child...
				self.parent.numChildren -= 1
				self.parent.prune()

	# ------------------
	def collapseUp(self):
	# ------------------
	# 'Collapse' the tree recursively by having nodes with only one child get cut out and 'absorbed' by the child
		# Collapse this node up as far as it'll go
		while self.parent is not None and self.parent.numChildren == 1:
			# Create a set out of all of our updates 
			
			# 'Own' all the parents updates
			for pt in self.parent.updated:
				# Alias the grid cell and remove the parents value
				cell = self.slamMap[pt]
				val = cell[self.parent.particleId]
				del cell[self.parent.particleId]

				# Own the points you don't already update yourself
				if pt not in self.updated:
					# Change the point's owner in the map cell
					cell[self.particleId] = val

					# Add the point to our updated list	
					self.updated.add(pt)

				# If you change the point's value right back, you can remove it from your update list
				# (the successive changes cancel each other out)
				else:
					del cell[self.particleId]
					self.updated.remove(pt)
				
			# Splice out the parent
			self.parent = self.parent.parent

		# Recurse up the tree
		if self.parent is not None:
			self.parent.collapseUp()

	# ------------------------------
	def motionModel(self, deltaPose):
	# ------------------------------
		# Otherwise, update pose based on deltaPose
		"""
		self.position = (self.position[0] + gauss(deltaPose[0], deltaPose[0] * MOTION_SIGMA), \
						 self.position[1] + gauss(deltaPose[1], deltaPose[1] * MOTION_SIGMA))
		"""

		# Calcuate the distance
		distance = sqrt(pow(deltaPose[0],2) + pow(deltaPose[1],2))
		distance = gauss(distance, distance * MOTION_SIGMA)

		# Calculate the new heading
		heading = self.heading
		newHeading = heading + gauss(deltaPose[2], ROTATION_SIGMA)

		# Average the two (for position update)
		avgHeading = (newHeading + heading) / 2.0

		# Update the position
		self.position = (self.position[0] + (distance * cos(avgHeading)), self.position[1] + (distance * sin(avgHeading))) 
		
		# Update the heading
		self.heading = newHeading

	# ------------------------------------
	def getMapCell(self, pt, ancestryList):
	# ------------------------------------
		""" Reconstructs the map cell at the specified point """
		
		# Reconstruct the map cell at the point represented by the key
		cell = self.slamMap[pt]

		# The cell value defaults to unknown
		cellValue = MAP_CELL_UNKNOWN

		# Look for the first ancestor to have changed the value
		for ancestor in ancestryList:
			# If the ancestor changed the value of the cell, we're done
			if ancestor in cell:
				cellValue = cell[ancestor]
				break

		return cellValue

	# ----------------
	def __str__(self):
	# ----------------
		return 'p%i @ probability %f' % (self.particleId, self.probability)
